#!/usr/bin/env python

#################################################################################
#
# Copyright (c) 2006 Michigan State University
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#################################################################################

#################################################################################
#
#	Author:
#		Nathan Collins <npcollins@gmail.com>
#
#################################################################################

""" SRU interface for OAIS.

Provides an limited Search/Retrieval via URL (SRU) interface for OAIS.
Based off: http://www.loc.gov/standards/sru/index.html
"""

__version__ = "0.1"

# imports
import cgi
import PyLucene
from xml.sax import saxutils

# local imports
import cql
import fieldmapping
from oaisimilarity import oaiSimilarity

class xmlRoot:
	"""
	Base of XML tree for output of SRU response.
	"""
	def __init__(self):
		"""
		Initializes the root of the tree to None.
		"""
		self.root = None
		self.stylesheet = None

	def setRoot(self, newRoot):
		"""
		Sets the root of the tree to newRoot.
		"""
		self.root = newRoot

	def setStylesheet(self, stylesheet):
		"""
		Sets the stylesheet.
		"""
		self.stylesheet = stylesheet

class xmlNode:
	"""
	A node in the XML tree used for output of SRU response.
	"""
	def __init__(self):
		"""
		Class initialization, sets variables to None or empty.
		"""
		# name of tag. (e.g. srw:record)
		self.tag = None
		# attribute dict (e.g. {'xmlns:srw':'http://url/'})
		self.attr = {}
		# content between open and close tags
		self.content = None
		# children nodes
		self.children = []
		# parent node
		self.parent = None

	def setTag(self, tag, attrList = {}, content = None):
		self.tag = tag
		self.attr = attrList
		self.content = content

	def setChildren(self, children):
		self.children = children

	def addChild(self, child):
		self.children.append(child)

	def setParent(self, parent):
		self.parent = parent

def formXML(parsed):
	"""
	Create the XML tree from the parsed parameters.
	"""
	# create tree root
	tree = xmlRoot()

	### STYLESHEET
	stylesheet = ''
	# set stylesheet
	if len(parsed['stylesheet']) > 0:
		stylesheet = u'<?xml-stylesheet type="text/xsl" href="%s"?>' % parsed['stylesheet']

	# add to tree
	tree.setStylesheet(stylesheet)

	# ROOT
	searchRetrieveResponse = xmlNode()
	searchRetrieveResponse.setTag('searchRetrieveResponse', \
		{ 'xmlns:srw':'"http://www.loc.gov/zing/srw/"',
		'xmlns:diag':'"http://www.loc.gov/zing/srw/diagnostic/"',
		'xmlns:xcql':'"http://www.loc.gov/zing/cql/xcql/"',
		'xmlns:dc':'"http://purl.org/dc/elements/1.1/"' } )
	tree.setRoot(searchRetrieveResponse)

	# SEARCH FOR RESULTS
	results, total = getResults(parsed['lucene'], parsed['startRecord'], parsed['maximumRecords'])

	# RESULT INFO

	# version
	version = xmlNode()
	version.setTag('version', {}, str(parsed['version']) )
	tree.root.addChild(version)

	# numberOfRecords
	numberOfRecords = xmlNode()
	numberOfRecords.setTag('numberOfRecords', {}, str(total) )
	tree.root.addChild(numberOfRecords)

	# RECORDS

	# records
	records = xmlNode()
	records.setTag('records')

	# add results to records
	for result in results:
		# generate record node
		record = genRecordNode(result, parsed)
		records.addChild(record)

	# add to tree
	tree.root.addChild(records)

	# add diagnostic
	if len(parsed['sortKeys'].strip()) > 0:
		diag = diagnostic(80, 'Sorting is handled by Lucene.', 0)

		diagnode = xmlNode()
		diagnode.setTag('diagnostics', {}, diag)

		tree.root.addChild(diagnode)
		

	# return tree
	return tree


def genRecordNode(result, parsed):
	"""
	Given a results dict, generate a node to add to the tree.
	"""
	# create node
	record = xmlNode()
	record.setTag('record')

	# recordPacking
	recordPacking = xmlNode()
	recordPacking.setTag('recordPacking', {}, parsed['recordPacking'])
	record.addChild(recordPacking)

	# recordSchema
	recordSchema = xmlNode()
	recordSchema.setTag('recordSchema', {}, 'info:srw/schema/1/dc-v1.1')
	record.addChild(recordSchema)

	# recordData
	recordData = xmlNode()
	recordData.setTag('recordData')

	# dc.fields
	# dc
	dc = xmlNode()
	dc.setTag('dc:dc')

	# xml output
	for element in result:
		elemNode = xmlNode()
		# element in form of dc.element, modify for output
		tagelement = element
		if element[:3] == 'dc.': tagelement = 'dc:' + element[3:]

		elemNode.setTag(tagelement, {}, result[element])
		dc.addChild(elemNode)

	if parsed['recordPacking'] == 'xml':
		recordData.addChild(dc)
	
	elif parsed['recordPacking'] == 'string':
		content = '<?xml version="1.0">' + getChildXML(dc)
		content = saxutils.escape(content)
		content = content.encode('ascii','xmlcharrefreplace')

		# set record data string content
		recordData.setTag('recordData', {}, content)

	record.addChild(recordData)

	# recordNumber (optional, but recommended)

	# return node
	return record


def getResults(lucenequery, start = 1, max = 20):
	"""
	Search lucene and retrieve the the results. Return as a list of dicts.
	"""
	# lucene directory
	lucenedir = '/var/lib/lucene/harvest'

	# list of results
	results = []
	totalresults = 0

	# length checking
	if len(lucenequery) > 1:
		# the lucene searcher
		searcher = PyLucene.IndexSearcher(lucenedir)

		# set the similarity of the searcher
		searcher.setSimilarity(oaiSimilarity())

		# create lucene query parser, default search field is 'keywords'
		parser = PyLucene.QueryParser("keywords", PyLucene.StandardAnalyzer())

		# parse query
		readyquery = parser.parse(lucenequery)

		# perform search
		luceneresults = searcher.search(readyquery)

		# total possible results
		totalresults = len(luceneresults)

		# count for current position in loop
		current = 1

		# loop through lucene's results
		for i, doc in luceneresults:
			# check to see if we at a position where we want to pull results
			if current >= start and current < (start + max):
				# pull the record into a dict
				record = {}
				for field in doc:
					record[field.name()] = field.stringValue()

				# escape brackets and amp
				record = dict([(key, saxutils.escape(record[key])) for key in record])
				# pull dc dict from record dict
				dc = fieldmapping.recordmapactual2dc(record)

				# append to results list
				results.append(dc)
			
			# if we are past what is needed, break out
			if current >= (start + max):
				break

			# increment position
			current += 1

	# return results
	return results, totalresults

def getParams():
	"""
	Returns: dict
	
	Get CGI parameters for a SRU request and return them in a dictionary.
	"""
	fields = ['version','query','startRecord','maximumRecords','recordSchema',
		'recordPacking','recordXPath','resultSetTTL','sortKeys',
		'stylesheet','extraRequestData','operation']
	params = cgi.FieldStorage()
	paramDict = {}
	for field in fields:
		if params.has_key(field):
			paramDict[field] = params[field].value
		else:
			paramDict[field] = ''
	return paramDict

def parseParams(params):
	"""
	Parse the parameters sent as a request and determine the response to take.
	"""
	parsed = { 'lucene':'' }
	### version (mandatory)
	# version 1.0 was a trial version, servers are not required to respond with 1.0
	if params['version'] != '1.1':
		# if version other than 1.1 is requested
		# atm, responding with version 1.1 regardless
		parsed['version'] = '1.1'
	else:
		parsed['version'] = params['version']
	
	### query (mandatory)
	if params.has_key('query'):
		parsed['query'] = params['query']
		parsed['lucene'] = cql.toLucene(params['query'])
		# replace dc.fields with actual db fields
		parsed['lucene'] = fieldmapping.querymapdc2actual(parsed['lucene'])

	### startRecord

	# where to start results, lowest record starts at 1
	try:
		parsed['startRecord'] = abs(int(params['startRecord']))
		if parsed['startRecord'] > 1000: parsed['startRecord'] = 1
	except:
		parsed['startRecord'] = 1

	### maximumRecords

	# the number of records to return, can be less than requested, but never more
	try:
		parsed['maximumRecords'] = abs(int(params['maximumRecords']))
		if parsed['maximumRecords'] > 20: parsed['maximumRecords'] = 20
	except:
		parsed['maximumRecords'] = 10

	### recordPacking (xml or string)
	if params.has_key('recordPacking'):
		if params['recordPacking'] == 'xml' or params['recordPacking'] == 'string':
			parsed['recordPacking'] = params['recordPacking']
		else:
			parsed['recordPacking'] = 'xml'
	else:
		# defaults to xml
		parsed['recordPacking'] = 'xml'

	### recordSchema
	parsed['recordSchema'] = 'dc'	# currently only support dc

	### recordXPath
	parsed['recordXPath'] = params['recordXPath']

	### resultSetTTL

	### sortKeys
	parsed['sortKeys'] = params['sortKeys']

	### stylesheet (user provided)
	parsed['stylesheet'] = params['stylesheet']

	### extraRequestData

	### operation (mandatory)
	if not params.has_key('operation') or len(params['operation'].strip()) < 1: 
		params['operation'] = 'explain'
	
	parsed['operation'] = params['operation']

	# return parsed variables
	return parsed


def getChildXML(node):
	"""
	Generate and return the xml for the child node.
	"""
	tag = node.tag

	attr = ''

	for atts in node.attr:
		attr += ' %s=%s' % (atts, node.attr[atts])
	
	if node.content:
		content = node.content.encode('ascii','xmlcharrefreplace')
	else:
		content = ''
		for child in node.children:
			content += getChildXML(child) + '\n'

	# empty tag returns empty string
	if len(attr) + len(content) == 0:
		return ''
	
	return '<%s%s>%s</%s>' % (tag, attr, content, tag)


def createXML(tree):
	"""
	Generate xml for returning.
	"""

	# xml response string
	xml = "<?xml version='1.0'?>\n"

	# add stylesheet
	if len(tree.stylesheet) > 0:
		xml += tree.stylesheet + '\n'

	xml += getChildXML(tree.root)
	
	# return xml
	return xml


def explain(parsed, diag = ''):
	"""
	Return an explainResponse.
	"""
	### STYLESHEET
	stylesheet = ''
	# set stylesheet
	if len(parsed['stylesheet']) > 0:
		stylesheet = u'<?xml-stylesheet type="text/xsl" href="%s"?>' % parsed['stylesheet']

	# response
	host = 'lib0026.lib.msu.edu'
	port = '7777'
	database = 'oais'

	response = """<explainResponse>
		<version>1.1</version>
		<record>
		<recordSchema>http://explain.z3950.org/dtd/2.0/</recordSchema>
		<recordPacking>xml</recordPacking>
		<recordData>
		<explain>
		<serverInfo>
		<host>%s</host>
		<port>%s</port>
		<database>%s</database>
		</serverInfo>
		</explain>
		</recordData>
		</record>
		%s
		</explainResponse>
		""" % (host, port, database, diag)

	# return explainResponse
	return stylesheet + '\n' + response


def diagnostic(number, details, fatal = 1):
	"""
	Generate a diagnostics response.
	"""
	# list of diag types
	diagtype = { 	4: 'Unsupported operation.',
			5: 'Unsupported version.',
			80: 'Sort not supported.'
		}


	diag = """<diagnostic>
		<uri>info:srw/diagnostic/1/%s</uri>
		<message>%s</message>
		<details>%s</details>
		</diagnostic>
		""" % (number, diagtype[number], details)
	
	if fatal:
		diag = """<diagnostics>
			%s
			</diagnostics>
			""" % diag

	return diag


def main():
	"""
	Returns: None

	Main function. This is called upon execution.
	"""
	params = getParams()

	# page header
	print "Content-type: text/xml; charset=UTF-8\n"
	parsed = parseParams(params)

	# searchRetrieve process
	if parsed['operation'] == 'searchRetrieve':
		# form up the xml tree
		tree = formXML(parsed)

		# output tree to xml
		xml = createXML(tree)

	# explain process
	elif parsed['operation'] == 'explain':
		# form the explainResponse
		xml = explain(parsed)
	
	# otherwise unknown operation
	else:
		diag = diagnostic(4, parsed['operation'])

		xml = explain(parsed, diag)

	# return xml
	print xml


if __name__ == '__main__' or __name__ == 'sru':
	main()

# eof
